import math
import random

class PSO(object):
    def __init__(self):
        self.capacity = 120  # 车辆最大容量
        self.CostStartCar = 100  # 车辆启动成本
        self.CostDriving = 1  # 车辆单位距离行驶成本
        self.birdNum = 50  # 粒子数量
        self.w = 0.2  # 惯性因子
        self.c1 = 0.4  # 自我认知因子
        self.c2 = 0.4  # 社会认知因子
        self.iterMax = 5  # 迭代次数

    def greedy(self, CityCoordinates, dis_matrix):
        '''
        贪婪策略构造初始解,初始化时将VRP简化为TSP进行构造。
        输入：CityCoordinates-节点坐标,dis_matrix-距离矩阵
        输出：初始解-line
        '''
        # 修改dis_matrix以适应求解需要
        dis_matrix = dis_matrix.astype('float64')
        for i in range(len(CityCoordinates)):
            dis_matrix.loc[i, i] = math.pow(10, 10)

        dis_matrix.loc[:, 0] = math.pow(10, 10)  # 0不在编码内
        line = []  # 初始化
        now_city = random.randint(1, len(CityCoordinates) - 1)  # 随机生成出发站点
        line.append(now_city)  # 添加当前站点到路径
        dis_matrix.loc[:, now_city] = math.pow(10, 10)  # 更新距离矩阵，已经过站点不再被取出

        for i in range(1, len(CityCoordinates) - 1):
            next_city = dis_matrix.loc[now_city, :].idxmin()  # 距离最近的站点
            line.append(next_city)  # 添加进路径
            dis_matrix.loc[:, next_city] = math.pow(10, 10)  # 更新距离矩阵
            now_city = next_city  # 更新当前站点
        return line

    def calFitness(self, birdPop, Demand, dis_matrix, capacity, CostStartCar, CostDriving):
        '''
        贪婪策略分配车辆（解码），计算路径距离（评价函数）
        输入：birdPop-路径，Demand-客户需求,dis_matrix-站点间距离矩阵，capacity-车辆最大载重,DISTABCE-车辆最大行驶距离,CostStartCar-车辆启动成本,CostDriving-车辆单位距离行驶成本；
        输出：birdPop_car-分车后路径,fits-适应度
        '''
        birdPop_car, fits = [], []  # 初始化
        for j in range(len(birdPop)):
            bird = birdPop[j]
            lines = []  # 存储线路分车
            line = [0]  # 每辆车服务客户点
            dis_sum = 0  # 线路距离
            dis, d = 0, 0  # 当前客户距离前一个客户的距离、当前客户需求量
            i = 0  # 指向配送中心
            while i < len(bird):
                if line == [0]:  # 车辆未分配客户点
                    dis += dis_matrix.loc[0, bird[i]]  # 记录距离
                    line.append(bird[i])  # 为客户点分车
                    d += Demand[bird[i]]  # 记录需求量
                    i += 1  # 指向下一个客户点
                else:  # 已分配客户点则需判断车辆载重和行驶距离
                    if d + Demand[bird[i]] <= capacity:  #(dis_matrix.loc[line[-1], bird[i]] + dis_matrix.loc[bird[i], 0] + dis <= DISTABCE) &
                        dis += dis_matrix.loc[line[-1], bird[i]]
                        line.append(bird[i])
                        d += Demand[bird[i]]
                        i += 1
                    else:
                        dis += dis_matrix.loc[line[-1], 0]  # 当前车辆装满
                        line.append(0)
                        dis_sum += dis
                        lines.append(line)
                        # 下一辆车
                        dis, d = 0, 0
                        line = [0]
            # 最后一辆车
            dis += dis_matrix.loc[line[-1], 0]
            line.append(0)
            dis_sum += dis
            lines.append(line)
            birdPop_car.append(lines)
            fits.append(round(CostDriving * dis_sum + CostStartCar * len(lines), 1))
        return birdPop_car, fits

    def crossover(self, bird, pLine, gLine, w, c1, c2):
        '''
        采用顺序交叉方式；交叉的parent1为粒子本身，分别以w/(w+c1+c2),c1/(w+c1+c2),c2/(w+c1+c2)
        的概率接受粒子本身逆序、当前最优解、全局最优解作为parent2,只选择其中一个作为parent2；
        输入：bird-粒子,pLine-当前最优解,gLine-全局最优解,w-惯性因子,c1-自我认知因子,c2-社会认知因子；
        输出：交叉后的粒子-croBird；
        '''
        croBird = [None] * len(bird)  # 初始化
        parent1 = bird  # 选择parent1
        # 选择parent2（轮盘赌操作）
        randNum = random.uniform(0, sum([w, c1, c2]))
        if randNum <= w:
            parent2 = [bird[i] for i in range(len(bird) - 1, -1, -1)]  # bird的逆序
        elif randNum <= w + c1:
            parent2 = pLine
        else:
            parent2 = gLine
        # parent1-> croBird
        start_pos = random.randint(0, len(parent1) - 1)
        end_pos = random.randint(0, len(parent1) - 1)
        if start_pos > end_pos: start_pos, end_pos = end_pos, start_pos
        croBird[start_pos:end_pos + 1] = parent1[start_pos:end_pos + 1].copy()
        # parent2 -> croBird
        list2 = list(range(0, start_pos))
        list1 = list(range(end_pos + 1, len(parent2)))
        list_index = list1 + list2  # croBird从后往前填充
        j = -1
        for i in list_index:
            for j in range(j + 1, len(parent2) + 1):
                if parent2[j] not in croBird:
                    croBird[i] = parent2[j]
                    break
        return croBird

    def iteration(self, Customer, dis_matrix, Demand):
        birdPop = [self.greedy(Customer, dis_matrix) for i in range(self.birdNum)]  # 贪婪算法构造初始解

        birdPop_car, fits = self.calFitness(birdPop, Demand, dis_matrix,
                                            self.capacity,
                                            self.CostStartCar,
                                            self.CostDriving)  # 分配车辆，计算种群适应度

        gBest = min(fits)  # 全局最优值、当前最优值
        gLine = pLine = birdPop[fits.index(min(fits))]  # 全局最优解、当前最优解
        gLine_car = birdPop_car[fits.index(min(fits))]

        iterI = 1
        while iterI <= self.iterMax:  # 迭代开始
            for i in range(self.birdNum):
                birdPop[i] = self.crossover(birdPop[i], pLine, gLine, self.w, self.c1, self.c2)

            birdPop_car, fits = self.calFitness(birdPop, Demand, dis_matrix,
                                                self.capacity,
                                                self.CostStartCar,
                                                self.CostDriving)  # 分配车辆，计算种群适应度
            pBest, pLine, pLine_car = min(fits), birdPop[fits.index(min(fits))], birdPop_car[fits.index(min(fits))]
            if min(fits) <= gBest:
                gBest, gLine, gLine_car = min(fits), birdPop[fits.index(min(fits))], birdPop_car[fits.index(min(fits))]

            print(f"第{iterI}次迭代，最佳适应度：{gBest}")  # 打印当前代数和最佳适应度值
            iterI += 1  # 迭代计数加一
        return gLine_car

